{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hackathon 7\n",
    "\n",
    "Written by Eleanor Quint\n",
    "\n",
    "Topics:\n",
    "- Techniques for dimension expansion\n",
    "    - Transpose convolutions\n",
    "    - Sub-pixel convolutions\n",
    "    - ProgressiveGAN upscaling\n",
    "- Autoencoding\n",
    "    - Sparse autoencoders\n",
    "    - De-noising autoencoders\n",
    "\n",
    "This is all setup in a IPython notebook so you can run any code you want to experiment with. Feel free to edit any cell, or add some to run your own code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll start with our library imports...\n",
    "from __future__ import print_function\n",
    "\n",
    "import numpy as np                 # to use numpy arrays\n",
    "import tensorflow as tf            # to specify and run computation graphs\n",
    "import tensorflow_datasets as tfds # to load training data\n",
    "import matplotlib.pyplot as plt    # to visualize data and draw plots\n",
    "from tqdm import tqdm              # to track progress of loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's use the code from Hack2 to load MNIST\n",
    "ds = tfds.load('mnist', shuffle_files=True) # this loads a dict with the datasets\n",
    "\n",
    "# We can create an iterator from each dataset\n",
    "# This one iterates through the train data, shuffling and minibatching by 32\n",
    "train_ds = ds['train'].shuffle(1024).batch(32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Techniques for dimension expansion\n",
    "\n",
    "Generally, in the classifiers we've used so far, we compress high dimensional representations into lower dimensional ones. Now, we're going to study ways of going from lower dimensional to higher. For this, we're going to define a function `upscale_block` which makes the representation larger, in three different ways.\n",
    "\n",
    "#### Transpose convolutions\n",
    "\n",
    "Although we can upscale unstructured data very easily with a larger dense layer, images are less straightforward to upscale. As the upscaling companion to downscaling convolutions, we can learn to increase the size of images with transpose convolutions. They are sometimes called \"deconvolutions\" because they're the inverse operation of the convolution, but it is actually the transpose (gradient) of a convolution rather than an actual deconvolution. Transpose convolutions are implemented by [tf.keras.layers.Conv2DTranspose\n",
    "](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2DTranspose)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upscale_block(filters, kernel_size=3, scale=2, activation=tf.nn.elu):\n",
    "    \"\"\"conv2d_transpose\"\"\"\n",
    "    return tf.keras.layers.Conv2DTranspose(filters, kernel_size, strides=(scale, scale), padding='same', activation=activation)\n",
    "\n",
    "upscale_op = upscale_block(64)\n",
    "for batch in train_ds:\n",
    "    img = tf.cast(batch['image'], tf.float32)\n",
    "    up_x = upscale_op(img)\n",
    "    print(\"Original shape:\", img.shape, \"New shape:\", up_x.shape)\n",
    "    break\n",
    "\n",
    "num_params = upscale_op.count_params()\n",
    "print('Parameters: ' + str(num_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sub-pixel convolutions\n",
    "\n",
    "Another approach is called the sub-pixel convolution, which does a regular convolution with many channels, and then re-orders the data into the height and width dimensions from the channels dimension. For some intuition, a graphic:\n",
    "\n",
    "<img src=\"https://ai2-s2-public.s3.amazonaws.com/figures/2017-08-08/03a5b2aac53443e6078f0f63b35d4f95d6d54c5d/2-Figure1-1.png\">\n",
    "\n",
    "(Image sourced from [Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network](https://arxiv.org/abs/1609.05158))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upscale_block(filters, kernel_size=3, scale=2, activation=tf.nn.elu):\n",
    "    \"\"\"[Sub-Pixel Convolution](https://arxiv.org/abs/1609.05158)\"\"\"\n",
    "    # Increase the number of channels to the number of channels times the scale squared\n",
    "    conv = tf.keras.layers.Conv2D(filters * (scale ** 2), (kernel_size, kernel_size), activation=activation, padding='same')\n",
    "    # Rearrange blocks of (1,1,scale**2) pixels into (scale,scale,1) pixels\n",
    "    rearrange = tf.keras.layers.Lambda(lambda x: tf.nn.depth_to_space(x, scale))\n",
    "    return tf.keras.Sequential([conv, rearrange])\n",
    "    \n",
    "upscale_op = upscale_block(64)\n",
    "for batch in train_ds:\n",
    "    img = tf.cast(batch['image'], tf.float32)\n",
    "    up_x = upscale_op(img)\n",
    "    print(\"Original shape:\", img.shape, \"New shape:\", up_x.shape)\n",
    "    break\n",
    "\n",
    "num_params = upscale_op.count_params()\n",
    "print('Parameters: ' + str(num_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that this upscale method uses more parameters than the transpose convolution. This makes it more powerful and flexible (though this isn't always desirable, e.g., if your model is overfitting).\n",
    "\n",
    "#### ProgressiveGAN upscaling\n",
    "\n",
    "Finally, one technique that's recently found massive success in Nvidia's ProgressiveGAN used to generate high-resolution fake celebrity faces:\n",
    "\n",
    "<img src=\"https://i2.wp.com/robotnyheter.se/wp-content/uploads/2018/01/Nvidia_GAN_ansikten.jpg?w=1561\" width=\"70%\">\n",
    "\n",
    "None of these are real photos, they've all been upsampled in the GAN framework from Gaussian noise with an architecture that uses this technique. It upscales using a classical algorithm like bilinear upscaling and then transforms the output with convolutional layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upscale_block(filters, kernel_size=3, scale=2, activation=tf.nn.elu):\n",
    "    \"\"\" similar to the upsampling used in [ProgressiveGAN](https://arxiv.org/pdf/1710.10196.pdf) \"\"\"\n",
    "    def upscale(x):\n",
    "        n, w, h, c = x.get_shape().as_list()\n",
    "        return tf.image.resize(x, [scale*h, scale*w])\n",
    "    \n",
    "    upscale_layer = tf.keras.layers.Lambda(upscale)\n",
    "    conv = tf.keras.layers.Conv2D(filters, kernel_size, padding='same', activation=activation)\n",
    "    return tf.keras.Sequential([upscale_layer, conv])\n",
    "\n",
    "upscale_op = upscale_block(64)\n",
    "for batch in train_ds:\n",
    "    img = tf.cast(batch['image'], tf.float32)\n",
    "    up_x = upscale_op(img)\n",
    "    print(\"Original shape:\", img.shape, \"New shape:\", up_x.shape)\n",
    "    break\n",
    "\n",
    "num_params = upscale_op.count_params()\n",
    "print('Parameters: ' + str(num_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autoencoding\n",
    "\n",
    "Generally, autoencoding is learning \"a complicated identity function\". This makes it a form of unsupervised learning, which doesn't require data to be explicitly labeled, but instead looks for patterns and trends in data. Typically the complication is to bottleneck the size of the representation, but can also be more varied. We'll look at code for sparse autoencoders and de-noising autoencoders.\n",
    "\n",
    "First we'll define some preliminaries that we'll use in both architectures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upscale_block(filters, kernel_size=3, scale=2, activation=tf.nn.elu):\n",
    "    \"\"\"[Sub-Pixel Convolution](https://arxiv.org/abs/1609.05158)\"\"\"\n",
    "    # Increase the number of channels to the number of channels times the scale squared\n",
    "    conv = tf.keras.layers.Conv2D(filters * (scale**2),\n",
    "                                  (kernel_size, kernel_size),\n",
    "                                  activation=activation,\n",
    "                                  padding='same')\n",
    "    # Rearrange blocks of (1,1,scale**2) pixels into (scale,scale,1) pixels\n",
    "    rearrange = tf.keras.layers.Lambda(\n",
    "        lambda x: tf.nn.depth_to_space(x, scale))\n",
    "    return tf.keras.Sequential([conv, rearrange])\n",
    "\n",
    "\n",
    "class UpscaleBlock(tf.keras.layers.Layer):\n",
    "    def __init__(self, number, kernel_size=3, activation=tf.nn.swish):\n",
    "        super().__init__(name=\"UpscaleBlock\" + str(number))\n",
    "        self.activation = activation\n",
    "        self.kernel_size = kernel_size\n",
    "        self.is_built = False\n",
    "\n",
    "    def build(self, x):\n",
    "        channels = x.shape.as_list()[-1]\n",
    "        filters = channels // 2\n",
    "\n",
    "        bn1 = tf.keras.layers.BatchNormalization()\n",
    "        conv1 = upscale_block(filters)\n",
    "        bn2 = tf.keras.layers.BatchNormalization()\n",
    "        conv2 = tf.keras.layers.Conv2D(filters,\n",
    "                                       self.kernel_size,\n",
    "                                       padding='same')\n",
    "        self.main_network = [self.activation, bn1, conv1, self.activation, bn2, conv2]\n",
    "\n",
    "        self.skip_connection = upscale_block(filters)\n",
    "        self.se_activate = SqueezeExcite(filters)\n",
    "        self.is_built = True\n",
    "\n",
    "    def __call__(self, input_):\n",
    "        if not self.is_built:\n",
    "            self.build(input_)\n",
    "        x = input_\n",
    "        for layer in self.main_network:\n",
    "            x = layer(x)\n",
    "        output = x\n",
    "        skip = self.skip_connection(input_)\n",
    "        return skip + 0.1 * output\n",
    "\n",
    "\n",
    "class FactorizedReduce(tf.Module):\n",
    "    \"\"\"Downscale version of the sub-pixel convolution which re-arranges pixels\"\"\"\n",
    "    def __init__(self, channels):\n",
    "        super(FactorizedReduce, self).__init__()\n",
    "        assert channels % 2 == 0\n",
    "        self.conv_1 = tf.keras.layers.Conv2D(channels // 4, 1, strides=2)\n",
    "        self.conv_2 = tf.keras.layers.Conv2D(channels // 4, 1, strides=2)\n",
    "        self.conv_3 = tf.keras.layers.Conv2D(channels // 4, 1, strides=2)\n",
    "        self.conv_4 = tf.keras.layers.Conv2D(channels - 3 * (channels // 4),\n",
    "                                             1,\n",
    "                                             strides=2)\n",
    "        self.convs = [self.conv_1, self.conv_2, self.conv_3, self.conv_4]\n",
    "\n",
    "    def __call__(self, x):\n",
    "        \"\"\"Assumes NHCW data\"\"\"\n",
    "        assert x.shape[2] > 1\n",
    "        assert x.shape[3] > 1\n",
    "        out = tf.nn.swish(x)\n",
    "        conv1 = self.conv_1(out)\n",
    "        conv2 = self.conv_2(out[:, :, 1:, 1:])\n",
    "        conv3 = self.conv_3(out[:, :, :, 1:])\n",
    "        conv4 = self.conv_4(out[:, :, 1:, :])\n",
    "        out = tf.concat([conv1, conv2, conv3, conv4], -1)\n",
    "        return out\n",
    "\n",
    "\n",
    "class SqueezeExcite(tf.Module):\n",
    "    \"\"\"Activation function that performs gating\"\"\"\n",
    "    def __init__(self, out_channels):\n",
    "        super().__init__()\n",
    "        num_hidden = max(out_channels // 16, 4)\n",
    "        self.net = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(num_hidden), tf.keras.layers.Lambda(tf.nn.relu),\n",
    "            tf.keras.layers.Dense(out_channels), tf.keras.layers.Lambda(tf.nn.sigmoid)\n",
    "        ])\n",
    "\n",
    "    def __call__(self, x):\n",
    "        \"\"\"The choice of axes assumes we're working with NHWC data\"\"\"\n",
    "        ax = tf.math.reduce_mean(x, axis=[1, 2])\n",
    "        # data should be flat at this po,int\n",
    "        bx = self.net(ax)\n",
    "        cx = tf.expand_dims(tf.expand_dims(bx, 1), 1)\n",
    "        return cx * x\n",
    "\n",
    "\n",
    "class DownscaleBlock(tf.keras.layers.Layer):\n",
    "    def __init__(self, number, kernel_size=3, activation=tf.nn.swish):\n",
    "        super().__init__(name=\"DownscaleBlock\" + str(number))\n",
    "        self.activation = activation\n",
    "        self.kernel_size = kernel_size\n",
    "        self.is_built = False\n",
    "\n",
    "    def build(self, x):\n",
    "        channels = x.shape.as_list()[-1]\n",
    "        filters = channels * 2\n",
    "\n",
    "        bn1 = tf.keras.layers.BatchNormalization()\n",
    "        conv1 = tf.keras.layers.Conv2D(filters,\n",
    "                                       self.kernel_size,\n",
    "                                       strides=2,\n",
    "                                       padding='same')\n",
    "        bn2 = tf.keras.layers.BatchNormalization()\n",
    "        conv2 = tf.keras.layers.Conv2D(filters,\n",
    "                                       self.kernel_size,\n",
    "                                       padding='same')\n",
    "        self.main_network = [self.activation, bn1, conv1, self.activation, bn2, conv2]\n",
    "\n",
    "        self.skip_connection = FactorizedReduce(filters)\n",
    "        self.se_activate = SqueezeExcite(filters)\n",
    "        self.is_built = True\n",
    "\n",
    "    def __call__(self, input_):\n",
    "        if not self.is_built:\n",
    "            self.build(input_)\n",
    "        x = input_\n",
    "        for layer in self.main_network:\n",
    "            x = layer(x)\n",
    "        output = x\n",
    "        skip = self.skip_connection(input_)\n",
    "        return skip + 0.1 * output\n",
    "\n",
    "\n",
    "encoder_network = tf.keras.Sequential([\n",
    "    tf.keras.layers.Conv2D(32, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  #28,28,16\n",
    "    DownscaleBlock(1),  # 14,14,32\n",
    "    DownscaleBlock(2),  # 7,7,64\n",
    "    tf.keras.layers.Conv2D(64, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  # 7,7,64\n",
    "    tf.keras.layers.Conv2D(16, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  # 7,7,16\n",
    "    tf.keras.layers.Conv2D(1, 3, padding='same'),  # 7,7,4\n",
    "])\n",
    "\n",
    "decoder_network = tf.keras.Sequential([\n",
    "    tf.keras.layers.Conv2D(4, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  # 7,7,4\n",
    "    tf.keras.layers.Conv2D(16, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  # 7,7,16\n",
    "    tf.keras.layers.Conv2D(64, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  # 7,7,64\n",
    "    UpscaleBlock(1),  # 14,14,32\n",
    "    UpscaleBlock(2),  # 28,28,16\n",
    "    tf.keras.layers.Conv2D(4, 3, padding='same',\n",
    "                           activation=tf.nn.swish),  #28,28,4\n",
    "    tf.keras.layers.Conv2D(1, 3, padding='same'),  #28,28,16\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in train_ds:\n",
    "    x = tf.cast(batch['image'], tf.float32)\n",
    "    code = encoder_network(x)\n",
    "    output = decoder_network(code)\n",
    "    break\n",
    "print(x.shape, code.shape, output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sparse autoenoding\n",
    "\n",
    "Although we bottleneck the representation in normal autoencoding by reducing the dimensionality, sparse autoencoders can actually increase it, but restrict it to be sparsely activated with L1 regularization using [tf.norm](https://www.tensorflow.org/api_docs/python/tf/norm) or KL-divergence. This has the effect of only having non-zero values in a few dimensions, effectively bottlenecking each representation, but giving a greater variety of dimensions to choose to be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sparse_autoencoder_loss(x, code, x_hat, sparsity_coeff=5.):\n",
    "    sparsity_loss = tf.norm(code, ord=1, axis=1)\n",
    "    reconstruction_loss = tf.reduce_mean(tf.square(x_hat - x)) # Mean Square Error\n",
    "    total_loss = reconstruction_loss + sparsity_coeff * sparsity_loss\n",
    "    return total_loss\n",
    "\n",
    "max_steps = 250\n",
    "step = 0\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "for batch in tqdm(train_ds):\n",
    "    with tf.GradientTape() as tape:\n",
    "        x = tf.cast(batch['image'], tf.float32)\n",
    "        code = encoder_network(x)\n",
    "        output = decoder_network(code)\n",
    "        loss = sparse_autoencoder_loss(x, code, output)\n",
    "    gradient = tape.gradient(loss, encoder_network.trainable_variables + decoder_network.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradient, encoder_network.trainable_variables + decoder_network.trainable_variables))\n",
    "    step += 1\n",
    "    if step > max_steps:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show a code (which should be pretty sparse)\n",
    "print(code[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show reconstruction\n",
    "# This will perform much better with more training and hyperparameter tuning\n",
    "print(\"Original and Reconstruction\")\n",
    "side_by_side = tf.concat([x[0], output[0]], 1).numpy()\n",
    "plt.imshow(side_by_side)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Denoising Autoencoder\n",
    "\n",
    "Another way to force an autoencoder to learn the features of data is to force it to map noisy, corrupted versions of the data back to the original. This is usually accomplished by manually adding noice (e.g., Gaussian), but may also be useful in real world settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NOISE_COEFF = 10.\n",
    "\n",
    "def autoencoder_loss(x, x_hat):\n",
    "    reconstruction_loss = tf.reduce_mean(tf.square(x_hat - x)) # Mean Square Error\n",
    "    total_loss = reconstruction_loss\n",
    "    return total_loss\n",
    "\n",
    "max_steps = 250\n",
    "step = 0\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "for batch in tqdm(train_ds):\n",
    "    with tf.GradientTape() as tape:\n",
    "        x = tf.cast(batch['image'], tf.float32)\n",
    "        code = encoder_network(x + tf.random.normal(x.shape))\n",
    "        output = decoder_network(code)\n",
    "        loss = autoencoder_loss(x, output)\n",
    "    gradient = tape.gradient(loss, encoder_network.trainable_variables + decoder_network.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradient, encoder_network.trainable_variables + decoder_network.trainable_variables))\n",
    "    step += 1\n",
    "    if step > max_steps:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show a code (which should be pretty sparse)\n",
    "print(code[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show reconstruction\n",
    "# This will perform much better with more training and hyperparameter tuning\n",
    "print(\"Original and Reconstruction\")\n",
    "side_by_side = tf.concat([x[0] + NOISE_COEFF * tf.random.normal(x[0].shape), output[0]], 1).numpy()\n",
    "plt.imshow(side_by_side)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Homework\n",
    "\n",
    "Try out the autoencoder architecture above and test out the relationship between the mean squared error and the size of the latent variable (e.g., the (7,7,1) shape used above) using the above architecture on MNIST after some training. Try 2 or 3 code sizes (or more if you like) and report the parameters of a best fit line \n",
    "\n",
    "### Coda\n",
    "\n",
    "#### [Progressive GAN latent space interpolation on Youtube](https://youtu.be/XOxxPcy5Gr4?t=1m48s)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl_notebooks",
   "language": "python",
   "name": "dl_notebooks"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
